/*
 *  em.cpp
 *  bgsubtraction
 *
 *  Created by a1gucis on 3/17/09.
 *  Copyright 2009 __MyCompanyName__. All rights reserved.
 *
 */

#include "em.h"
/*
 *  emtest.cpp
 *  opencvTest
 *
 *  Created by a1gucis on 3/9/09.
 *  Copyright 2009 __MyCompanyName__. All rights reserved.
 *
 */

#include "em.h"
EM::EM() {
	samples = NULL;
	clusters = NULL;
	sampleIdx = 0;
	
	params.covs      = NULL;
    params.means     = NULL;
    params.weights   = NULL;
    params.probs     = NULL;
    params.nclusters = 4;
    params.cov_mat_type       = CvEM::COV_MAT_DIAGONAL;
    params.start_step         = CvEM::START_AUTO_STEP; // k means is used to estimate initial mixture parameters
    params.term_crit.max_iter = 10; // stop after a number of iterations
    params.term_crit.epsilon  = 0.01; // stop when the value change is very low
    params.term_crit.type     = CV_TERMCRIT_EPS;//CV_TERMCRIT_EPS | CV_TERMCRIT_ITER;//CV_TERMCRIT_ITER|CV_TERMCRIT_EPS; // use iterations as well as epsilon criterias.
	
}

EM::~EM() {
	cvReleaseMat(&samples);
	cvReleaseMat(&clusters);
}

void EM::setClusters(int clusters) {
	params.nclusters = clusters;
}

void EM::init(int height, int width) {
	if (samples != NULL)
		cvReleaseMat(&samples);
	if (clusters != NULL)
		cvReleaseMat(&clusters);
	samples = cvCreateMat(height*width,1,CV_32FC1); // matrix of all the scalar pixel of my src image 
	clusters = cvCreateMat(height*width,1,CV_32SC1);

}
void EM::train(IplImage *src, bool imageSizeDiffers) {
	if (imageSizeDiffers) {
		if (samples != NULL)
			cvReleaseMat(&samples);
		if (clusters != NULL)
			cvReleaseMat(&clusters);
		samples = cvCreateMat(src->height*src->width,1,CV_32FC1); // matrix of all the scalar pixel of my src image 
		clusters = cvCreateMat(src->height*src->width,1,CV_32SC1);
		
	}
	if (samples == NULL)
		samples = cvCreateMat(src->height*src->width, 1, CV_32FC1);
	if (clusters == NULL)
		clusters = cvCreateMat(src->height*src->width, 1, CV_32SC1);
	cvSetZero(clusters);
	//Utils::ToMatrix(src, samples, 0, 1);	
	//Utils::PrintDImg(src);
	Utils::ToDMatrix(src, samples);
	//Utils::PrintMat(samples);
	emModel.train(samples, 0, params, clusters);
	//cout<<"here 2"<<endl;
	cout<<"MEANS: "<<endl;
	const CvMat *means = emModel.get_means();
	//uchar *data = (uchar *)means->data.ptr;
	for (int i=0;i<means->rows;i++) {
		//double *data = (double *)(means->data.ptr + i * means->step);
		for (int j=0;j<means->cols;j++) {
			float mean = cvGetReal2D(means, i, j);
			cout<<mean<<"\t";
		}
		cout<<endl;
	}
	cout<<"##############"<<endl;
	
	
}

void EM::predict(IplImage *src, IplImage *labels) {
	CvMat *sample = cvCreateMat(1, 1, CV_32FC1);
	int srcStep = src->widthStep;
	int channels = src->nChannels;
	int labelsStep = labels->widthStep;
	uchar *srcData = (uchar *)src->imageData;
	uchar *labelsData = (uchar *)labels->imageData;
	for (int i=0;i<src->height;i++) {
		for (int j=0;j<src->width;j++) {
			if (channels == 1)
				cvSet1D(sample, 0, cvScalar(srcData[i*srcStep+j]));
			else {
				float chan1 = ((float *)(src->imageData + i*src->widthStep))[j*src->nChannels+0];
				float chan2 = ((float *)(src->imageData + i*src->widthStep))[j*src->nChannels+1];
				float chan3 = ((float *)(src->imageData + i*src->widthStep))[j*src->nChannels+2];
				float value = chan1 * W_CH1 + chan2 * W_CH2 + chan3 * W_CH3;
				//cvSetReal1D(sample, 0, value);
				cvSet1D(sample, 0, cvScalar(chan1 * W_CH1, chan2 * W_CH2, chan3 * W_CH3));
			}
				//cvSet1D(sample, 0, cvScalar(srcData[i*srcStep+j*channels+0], srcData[i*srcStep+j*channels+1], srcData[i*srcStep+j*channels+2]));
			int cluster = emModel.predict(sample, NULL);
			labelsData[i*labelsStep+j] = cluster;
		}
	}
	cvReleaseMat(&sample);
}

const CvMat* EM::get_means() {
	return emModel.get_means();
}

const CvMat** EM::get_covs() {
	return emModel.get_covs();
}

IplImage *EM::run2(IplImage *src, int numOfClusters) {
	IplImage* srcCopy = cvCreateImage(cvGetSize(src),IPL_DEPTH_8U, 1); 
	cvCopyImage(src,srcCopy); // create a copy to display 
	int nbElem = src->height * src->width;
	CvMat* samples = cvCreateMat(nbElem,1,CV_32FC1); // matrix of all the scalar pixel of my src image 
	CvMat* clusters = cvCreateMat(nbElem,1,CV_32SC1);
	cvSetZero(clusters);
	Utils::ToMatrix(src, samples);
	CvEM em_model;
    CvEMParams params;
	
	params.covs      = NULL;
    params.means     = NULL;
    params.weights   = NULL;
    params.probs     = NULL;
    params.nclusters = 2;
    params.cov_mat_type       = CvEM::COV_MAT_SPHERICAL;
    params.start_step         = CvEM::START_AUTO_STEP; // k means is used to estimate initial mixture parameters
    params.term_crit.max_iter = 10; // stop after a number of iterations
    params.term_crit.epsilon  = 0.1; // stop when the value change is very low
    params.term_crit.type     = CV_TERMCRIT_EPS;//CV_TERMCRIT_ITER|CV_TERMCRIT_EPS; // use iterations as well as epsilon criterias.
	
	time_t startTime = time(NULL);
	em_model.train(samples, 0, params, clusters);
	time_t endTime = time(NULL);
	cout<<"training took: "<<endTime-startTime<<" seconds"<<endl;
	//em_model.predict(samples, NULL);
	
	startTime = time(NULL);
	CvMat *sample = cvCreateMat(1, 1, CV_32FC1);
	
	int step = srcCopy->widthStep;
	for (int i=0;i<srcCopy->height;i++) {
		uchar *data = (uchar *)srcCopy->imageData;
		for (int j=0;j<srcCopy->width;j++) {
			cvSet1D(sample, 0, cvScalar(data[i*step+j]));
			int cluster = em_model.predict(sample, NULL);
			if (cluster == 0)
				data[i*step+j] = 0;
			else
				data[i*step+j] = 255;
		}
	}
	endTime = time(NULL);
	cout<<"predictions took: "<<endTime - startTime<<" seconds"<<endl;
	//cvSet1D(sample, 0, cvScalar(180));
	//cout<<em_model.predict(sample, NULL)<<endl;
	return srcCopy;
	
}
